from sqlalchemy.ext.asyncio import AsyncSession
from app.db.session import postgres_engine as engine, Base, PostgresSessionLocal as AsyncSessionLocal
from app.models.task import Task, TaskStatus
from app.models.user import User
import logging
import datetime
import traceback

logger = logging.getLogger(__name__)

async def init_db():
    """初始化数据库"""
    try:
        # 创建所有表
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
        
        # 添加初始数据
        await create_initial_data()
        
        logger.info("数据库初始化成功")
    except Exception as e:
        logger.error(f"数据库初始化失败: {str(e)}")
        logger.error(traceback.format_exc())
        raise

async def create_initial_data():
    """创建初始数据"""
    
    async with AsyncSessionLocal() as db:
        try:
            # 检查是否已有任务
            from sqlalchemy import select, func
            result = await db.execute(select(func.count()).select_from(Task))
            count = result.scalar()
            
            logger.info(f"当前任务数量: {count}")
            
            if count == 0:
                # 添加示例任务
                tasks = [
                    Task(
                        title="完成项目文档",
                        description="编写项目需求文档和设计文档",
                        status=TaskStatus.PENDING,
                        created_at=datetime.datetime.now()
                    ),
                    Task(
                        title="实现用户认证功能",
                        description="添加登录、注册和权限控制功能",
                        status=TaskStatus.IN_PROGRESS,
                        created_at=datetime.datetime.now()
                    ),
                    Task(
                        title="部署测试环境",
                        description="在测试服务器上部署应用",
                        status=TaskStatus.COMPLETED,
                        created_at=datetime.datetime.now()
                    ),
                    Task(
                        title="修复已知bug",
                        description="解决用户反馈的问题",
                        status=TaskStatus.FAILED,
                        created_at=datetime.datetime.now()
                    )
                ]
                
                for task in tasks:
                    db.add(task)
                
                await db.commit()
                logger.info(f"添加了 {len(tasks)} 个示例任务")
                
                # 验证数据是否添加成功
                result = await db.execute(select(func.count()).select_from(Task))
                new_count = result.scalar()
                logger.info(f"添加后任务数量: {new_count}")
        except Exception as e:
            logger.error(f"添加初始数据失败: {str(e)}")
            logger.error(traceback.format_exc())
            await db.rollback()
            raise

async def drop_db():
    """删除所有表"""
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.drop_all) 